The program below models a particle wavefunction dynamics in the field of an arbitrary time-independent potential. In particular, tunneling through a square barrier of finite height and width is demonstrated, as well as particle evolution in a harmonic potential. The method is based on solving numerically the time-independent Schrödinger equation. The results are visualized via animated plots.
The solution is based on solving (numerically) the time-independent Schrodinger equation for a given time-independent potential energy V: $$ \hat{H} \varphi = E\varphi \quad (1) $$ where $$ \hat{H} = -\frac{\hbar}{2m}\frac{d^{2}}{dx^{2}} + V(x) \quad (2) $$
Equation (1) represents an eigenvalue/eigenvector problem. The solution is a set of stationary basis wavefunctions (eigenfunctions), $\varphi_n$, each with the respective definite energy (eigenvalue) $E_n$.
The initial particle's wavefunction is represented as a moving wavepacket, $\psi(x,t=0)$. The set of the basis functions $\varphi_i$ is complete and orthonormal, and hence this wavepacket can be represented as their superposition:
$$
\psi(x,t=0) = \sum_n c_n \varphi_n (x)
$$
The wavefunction is uniquely described by vector $\mathbf c = (c_1 , c_2 ,c_3, ...)$ in the basis of eigenfunctions of the Hamiltonian.
The coefficients $c_n$ ('projections' to basis wavefunctions) can be easily found using a Fourier trick:
$$
c_n = \left < \varphi_n| \psi \right>
$$
Once the coefficients $c_n$ are known, the wavefunction at any later time can be expressed as the superposition of products of the time-independent $\varphi_n$ with its time-dependent counterpart $\exp{(-i E_n t/\hbar)}$:
$$\psi(x,t) = \sum_n c_n \varphi_n (x) \exp{\left ( -i E_n t/\hbar \right )}$$
In a numerical simulation, the operators and wavefunctions will be represented by matrices and vectors, respectively. Solving the problem will thus boil down to matrix multiplication.
To simplify the treatment, we will switch to atomic units where $\hbar = 1$ and we will assume that the particle's mass is that of an electron, $m = m_{electron} = 1$.
Note - a complete minimalistic version of the working Python program that animates propagation of a wavepacket through a barier is at the end of this document. You can run that last cell as a stand-alone routine; it contains copies of all necessary imports and functions.
In numerical modeling, the wavefunction is represented as a discrete function:
$$\psi(x,t) = \psi(x_i ,t) \equiv \psi _i$$
where x is an array of x-positions $\{x_i\}$ at which the function array $\{\psi _i\}$ is defined or calculated.
Assuming $\{x_i\}$ is a vector of size N:
$$\mathbf x = \left [ x_1, x_2, ..., x_N \right ] $$
The wavefunction is also a vector whose size is defined by the number of points in the array $\{x_i\}$:
$$
\boldsymbol \psi(x,t) = \left [ \psi(x_1 ,t), \psi(x_2 ,t), ..., \psi(x_N ,t) \right ] \equiv \left [\psi_1, \psi_2, ..., \psi_N \right ]
$$
Note that care should be taken in numerical multiplications of that vector with other vectors or matrices; in some cases, the vector should be a row, in others it should be a column.
The time-independent Schrödinger equation involves the product of the potential energy with the wavefunction:
$$
\hat{H} \varphi(x) = \frac{\hbar}{2m}\frac{d^{2}\varphi(x)}{dx^{2}} + V(x)\varphi(x)
$$
Since we work in a discrete space of $\{ x_i \}$, the product $V(x)\varphi(x) = V(x_i)\varphi(x_i) \equiv V_i\varphi_i$.
The potential function is, thus, also represented as a set of values $\{ V_i \}$ at positions $\{x_i\}$. For computational efficiency, we will represent the potential energy operator as a diagonal matrix $\mathbf V$ with its diagonal values being $\mathbf V_{ii} \equiv V(x_i) $. The product $V_i\varphi_i$ can then be represented as the following matrix product:
$$
V(x) \psi (x) = \mathbf{V}\boldsymbol \varphi = \begin{pmatrix}
&V(x_1) &0 &0 &... \\
&0 &V(x_2) &0 &... \\
&0 &0 &V(x_3) &... \\
&... &... &... &... \\
\end{pmatrix}
\begin{pmatrix}
\varphi(x_1) \\
\varphi(x_2) \\
\varphi(x_3) \\
... \\
\end{pmatrix}
= \begin{pmatrix}
V_1\varphi_1 \\
V_2\varphi_2 \\
V_3\varphi_3 \\
... \\
\end{pmatrix}
$$
The above matrix product multiplies the wavefunction vector by the potential at each point - the goal is achieved.
The right side of the time-independent Schrödinger equation $E \psi (x)$ involves the product of a constant by the wavefunction vector, and can also be represented as a matrix product: $$ E \psi (x) = \left( E\mathbf I \right ) \boldsymbol \varphi = \begin{pmatrix} &E &0 &0 &... \\ &0 &E &0 &... \\ &0 &0 &E &... \\ &... &... &... &... \\ \end{pmatrix} \begin{pmatrix} \varphi(x_1) \\ \varphi(x_2) \\ \varphi(x_3) \\ ... \\ \end{pmatrix} = \begin{pmatrix} E \varphi_1 \\ E \varphi_2 \\ E \varphi_3 \\ ... \\ \end{pmatrix} $$ where $\mathbf I$ is the identity matrix (unit matrix), i.e., the diagonal matrix with its diagonal elements being equal to 1. The above expression multiplies all components of vector $\boldsymbol \varphi$ by a scalar $E$. We will not need the last equation, it is presented here only for completeness of the discussion.
The Hamiltonian consists of two parts, kinetic and potential energies: $$ \hat{H} = \frac{\hbar}{2m}\frac{d^{2}}{dx^{2}} + V(x) $$ The first term is the kinetic energy $\hat T$, and it involves the second derivative of a numerical function $\psi (x)$ over x: $$ \hat{T} = \frac{\hbar}{2m}\frac{d^{2}}{dx^{2}} $$
Since we will be dealing with the wavefunction in numerical form that is known at discrete $x_i$ positions, we will express the derivative based on these values using the Finite Differences Method.
Denoting the grid interval between nearby points in $\{ x_i \}$ as $\Delta x$, we can approximate the first derivative at the midpoint position between $x_i$ and $x_{i+1}$ as:
Similarly, the first derivative at the position to the left of $x_i$ is:
$$ \frac{d\varphi(x_i - \Delta x/2)}{dx} \approx \frac {\varphi_{i} - \varphi_{i-1} }{\Delta x}$$Consequently, the second derivative is the difference between these two first derivatives over the x-distance between them:
$$ \frac{d^2 \varphi(x_i)}{dx^2} \approx \frac{\Delta \frac{d\varphi}{dx}}{\Delta x}= \frac{\frac{d\varphi(x_i + \Delta x/2)}{dx} - \frac{d\varphi(x_i - \Delta x/2)}{dx}}{\Delta x} = \frac {\frac {\varphi_{i+1} - \varphi_i }{\Delta x} - \frac {\varphi_{i} - \varphi_{i-1} }{\Delta x}}{\Delta x} $$This can be rewritten as: $$ \frac{d^2 \varphi(x_i)}{dx^2} \approx \frac {\varphi_{i-1} - 2 \varphi_i + \varphi_{i+1} }{\Delta x^2} $$
This result can be achieved by the following matrix multiplication: $$ \frac{d^{2}}{dx^{2}} \approx\frac {1}{\Delta x^2}\begin{pmatrix} &-2 &1 &0 &0 &0 &... \\ &1 &-2 &1 &0 &0 &... \\ &0 &1 &-2 &1 &0 &...\\ &0 &0 &1 &-2 &1 &... \\ &... &... &... &... &... &... \\ \end{pmatrix} \begin{pmatrix} \varphi_1 \\ \varphi_2 \\ \varphi_3 \\ \varphi_4 \\ ... \\ \end{pmatrix} \equiv \mathbf {D^{(2)}} \boldsymbol \varphi $$
where $\mathbf{D^{(2)}}$ denotes the second derivative matrix (including $1/\Delta x^2$). Indeed, for example, the second derivative at position $x_3$ is the third element of the resulting matrix multiplication, i.e., row 3 of the derivative matrix multiplied by the (first) column of the $\boldsymbol \varphi$-matrix, and we get:
$$ (\mathbf{D^{(2)}} \boldsymbol \varphi)_3 = \frac{1}{\Delta x^{2}} (1\varphi_2 - 2\varphi_3 +1\varphi_4) $$It is thus confirmed that multiplying $\mathbf{D^{(2)}}$ matrix by the wavefunction transforms the original wavefunction vector into its second derivative.
Note that the derivatives for the first and the last point are only correct if the function preceding the first point and the one following the last point are zeroes, so we should take care that the wavefunction is confined to within the x-range. If it extends beyond this region, the wavefunction will 'feel' the ends of the x-region as infinite walls and reflect from the edges.
One can also make the x-space periodic, i.e., point left of the first point, $x_1$ will be $x_N$, i.e., so that when the wavefunction runs beyond the edge of the x-space, it will appear on the other side:
$$
\frac{d^2 \varphi(x_1)}{dx^2} \approx \frac {\varphi_N - 2 \varphi_1 + \varphi_2 }{\Delta x^2}
$$
This is achieved by adding 1's to the top-right and bottom-left corners of the matrix $\mathbf {D^{(2)}}$. This will make the $\{x_i\}$-range periodic, i.e., the x-range is appended indefinitely to each side of the original x-array
Also note that the finer the step along the x-axis in the x-array, the more precise this approximation is.
We can, thus, represent the kinetic energy operator as the following matrix:
$$-\frac{\hbar}{2m}\frac{d^{2}}{dx^{2}}\varphi(x) \approx -\frac{\hbar}{2m} \frac {1}{\Delta x^2} \begin{pmatrix} &-2 &1 &0 &0 &0 &... &... &(1) \\ &1 &-2 &1 &0 &0 &... &... &0 \\ &0 &1 &-2 &1 &0 &... &... &0\\ &0 &0 &1 &-2 &1 &... &...&0 \\ &... &... &... &... &... &... &... &...\\ &... &... &... &... &... &... &... &1\\ &(1) &0 &0 &0 &0 &... &1 &-2 \end{pmatrix} \begin{pmatrix} \varphi_1 \\ \varphi_2 \\ \varphi_3 \\ \varphi_4 \\ ... \\ \end{pmatrix} = \mathbf{T} \boldsymbol \varphi$$where the kinetic energy operator matrix is $$ \mathbf T \equiv -\frac{\hbar}{2m} \mathbf {D^{(2)}} $$
The second derivative matrix above involves 1 point on each side of the center point, i.e., 3 points total - the minimum required for the second derivative, and it uses (1,-2,1) diagonal sequence. Higher accuracy can be achieved by taking into account more points around the center point. For example, a higher order of accuracy is achieved by diagonal sequence (-1/12, 4/3, -5/2, 4/3, -1/12): it uses 5 points (2 on each side of the center point)to obtain the second order derivative (see Finite difference coefficient)
Once the Hamiltonian matrix is known, the time-independent Schrödinger equation can be written and solved in vector form: $$ \mathbf H \boldsymbol \varphi = E \boldsymbol \varphi $$ It states that the matrix transformation of a vector $\boldsymbol \varphi$ is identical to multiplication by a constant (matrix diagonalization problem). Python has a number of functions that solve that equation for eigenvalues $E_n$ and eigenfunctions $\varphi_n$. We will be using numpy.linalg.eigh() function in this example.
Solving the time-independent Schrödinger equation yields a set of eigenvalues $E_i$ and eigenvectors (or eigenfunctions) $\varphi_i$. Each of these eigenfunctions represents a stationary wavefunction for which its temporal behavior is defined by the respective eigenvalue $E_i$. The task is to express the initial wavefunction $\psi(x,t=0)$ as a superposition of these eigenfunctions: $$ \psi(x,t=0) = \sum_i c_i \varphi_i (x) $$ The set of these eigenfunctions is orthonormal, i.e., the integrals: $$ \left<\varphi_i(x) | \varphi_j(x) \right> \equiv \int {\varphi_i(x_i)\varphi_j(x_i)dx} = \delta_{ij} $$ It provides an easy way of extracting the expansion coefficients $c_n$ using the following integral ("Fourier trick"): $$ c_n = \left<\varphi_n(x) | \psi(x) \right> $$ Indeed: $$ \left<\varphi_n | \psi \right> = \left<\varphi_n | \sum_i {c_i \varphi_i} \right> = \sum_i {c_i \left<\varphi_n |\varphi_i) \right>} = \sum_i {c_i \delta_{ni}} = c_n $$
In terms of discrete wavefunction vectors: $$ c_n = \sum_i \varphi_n(x_i) \psi(x_i) \Delta x= \sum_i \varphi_{ni}\psi_i \Delta x $$ or $$ \mathbf c = \boldsymbol {\varphi \psi} \Delta x $$
It is somewhat similar to extracting the projections of a geometrical vector onto one of the coordinates. For example, to obtain the projection to the x-axis of vector $\vec p = (x,y,z)$, one multiplies it with the unit vector $\hat i$ in the direction of the x-axis: $\vec p \hat i = (x\hat i + y \hat j + z \hat k)\hat i = x$. In analogy with geometrical vectors, where coordinates are projections to orthogonal axes, the values $c_n$ represent the "projections" of our function $\psi$ to the orthogonal basis functions $\varphi_n$.
Once we have the eigenfunctions $\varphi_n$, eigenvalues (energies) $E_n$, and the expansion coefficients $c_n$ for the initial wavefunction $\psi(x,t=0)$, we can find the wavefunction at any later time by appending the time-dependent part to each of the eigenfunctions in the expansion: $$ \psi(x_i,t) = \sum_n c_n \varphi_n (x_i) \exp{\left ( -i E_n t/\hbar \right )} \\ $$ Or, in matrix form: $$ \boldsymbol \psi(t) = \mathbf c \boldsymbol \varphi^T \left [ \exp (-i\frac{t}{\hbar} \mathbf E) \right ] $$ where $\boldsymbol \psi(t)$, $\mathbf E$, and $\mathbf c$ are vectors, the $\boldsymbol \varphi$ is a matrix such that $\boldsymbol \varphi_{ni} = \varphi_n(x_i) $. The $\exp()$ function exponents all elements of the matrix $\mathbf E$.
The import statement tells Python to load a library of functions. The statement can be followed by an optional as keyword that allows one to substitute the name of the original library to simplify its use.
The from statement allows importing a specific function from a library; once imported these functions can be called by their names without providing a full path. The * will import all functions from the library without the need to use the full path.
Finally, we set $\hbar =1$ and $m=1$, i.e., we will work in atomic units and assume that the mass of the particle is that of an electron.
save_mp4: set to False if you do not have ffmpeg installed separately and its bin added to system's PATH. ffmpeg is used to create mp4 animation file, the program will fail if that is not the case
show_html: Set to False if you only want to get mp4 movie. If True, shows the animation in browser, assuming that you run under Jupyter in browser. May fail in stand-alone Python that does not output to browser (never tested)
use_latex: When True animations will have equations written on them, but it will use an external LaTeX package. The program will fail if it is not the case, so set that to False. You can install package such as MiKTeX and provide system PATH to its bin directory
# By Sergei Savikhin, Purdue University, September 2025
import numpy as np # 'numerical python' library, defines matrices, comlpex numbers etc, we will refer to its functions using prefix 'np'
import matplotlib.pyplot as plt # set of functions for drawing graphs
from matplotlib.animation import FuncAnimation # set of functions to create animations using set of plots
from IPython.display import HTML,display # to export instructions in html format (inline Jupiter output)
from math import * # input all math functions (sin, cos, exp...), with no need to use any prefix
hbar = 1 # set hbar to 1
m = 1 # set mass to 1
save_mp4 = True # Save animation as MP4 file. Requires ffmpeg codec to be installed separately
# and its bin directory must be added to the system PATH to be found by Python
show_html = True # embed animation into output brtowser window. Will work under Jupyter
# but probably not in a stand alone Python like Spyder
use_latex = True # add symbols and equations to animations using external LaTex executables
# LaTeX packet (lieke MiKTeX) must be installed and path provided to its bin directory
The following 'helper' routine allows us to test functions by plotting them. It is not used in the animation routine.
The functions $y=y(x)$ are represented by x and y arrays, where for each $x_i$ there is a respective $y_i = y(x_i )$
Input parameters:
x - an array of x -values of the function
y - a list of y-arrays of one or more functions to plot. If y1 is an array of the first function, and y2 of the second, then this parameter is [y1,y2]. In case of just one function to plot, one still has to define a list with a single array in it [y1].
figsize = the width and height of the plot, the default is 10x6.
colors = an array of colors to be used for each y-graph, for example, colors = ['red','blue'] means that the first (y1) is red, and the second (y2) is blue
kwarg = additional parameters that will be passed to plot, for example, lw=0.5 will set the linewidth to 0.5.
The function returns the list of references to lines (can be ignored)
def Plot(x,y,size=(10,6),colors=[],**kwargs):
fig = plt.figure(figsize=size,dpi=100) # open figure of the defined size and dots per inch resolution
ax = plt.axes() #xlim=limits[0], ylim=limits[1]) #add axes, no limits - will autoadjust
for i in range(len(y)): # y may vave more than one curve, plot all
if i<len(colors): col = colors[i] # if colors array is defined - use specified color
else: col = f'C{i}' # if not set to default sequence C0, C1 .... f'' is a formatted text, it will print it as text but if a variable name is in {} its value will be printed
ax.plot(x,y[i],col,**kwargs) # append returned line references to list c
ax.grid() # plot default grid
#fig.show() # can show it, but usually it will show without this command as soon as the program is finished
return
We will generate the initial wavefunction $\psi(x,t=0)$ as a complex Gaussian wavepacket: $$ \psi(x,t=0) = A \exp{\left [-\frac{(x - x_0 )^2}{2 \sigma ^2}\right ]} \exp{(ik_0x)} $$ Here, the first term is the Gaussian function that defines the amplitude of the wave; its width is defined by $\sigma$, and the maximum position is at $x_0$. The second term is a complex oscillating function with wavenumber $k_0$. The wavefunction will be normalized by picking a correct A so that the integral of the $|\psi|^2$ over the whole space of x gives exactly 1: $$ \left < \psi(x) | \psi(x) \right > = \int {\left | \psi \right |^2dx} \approx \sum_i {\left | \psi_i \right |^2 \Delta x} = 1 $$ There is a connection between the k-number and the kinetic energy of the particle wave, $E_{wave} = \frac {p^2}{2m} = \frac {(\hbar k)^2}{2m}$. Therefore, a particle's wavepacket can be defined by either center wavevector $k_0$ as in the equation above, or by its kinetic energy $E_{particle}$. The expectation energy of the particle, $\left<E\right>$, described by the wavepacket that is symmetric along x, however, is not exactly equal to $(\hbar k)^2 /2m$ since the energy increases with k. One can obtain $k_0$ for the particle's expectation energy $\left<E\right>$ as:
$$ k_0 = \frac {m}{\hbar}\sqrt {2 \left < E \right >^2 - \frac {1}{2 \sigma^2}} $$Note that the kinetic energy of the particle could be found using: $$ \left < T \right > = \left < \psi | T | \psi \right > = \boldsymbol \psi^T \mathbf T \boldsymbol \psi $$ Since matrix calculations run in discrete space, the energy calculated using the above equation and respective matrices may yield slightly different results compared to the expectation energy used to set $k_0$.
Input parameters:
x = is an array of x-coordinates $\{x_i\}$ at which the wavefunction is defined. If no parameter is provided, it is set to range from 0 to 1000 and step $\Delta x$ =1
x0 = the position of the center of the wavepacket along x, i.e. $x_0$, if not provided the default value 200 is used
k0 - wavenumber of the central wave component (default 0.4)
sigma = $\sigma$, the width of the Gaussian (default is 15)
Note: for the second function, one defines the expectation energy of the particle E instead of k0
Output:
$\psi(x,t=0)$ array, $\{\psi_i\}$
Notes:
def Gaussian_wave(x = np.arange(0,1000), x0 = 200, k0 = 0.4,sigma = 15):
wf = np.exp( -1/2* (x-x0)**2/sigma**2) * np.exp(1j*k0*x)
wf = wf/(sqrt(np.sum(np.abs(wf)**2)*(x[1]-x[0]))) # normalize
return wf
def Gaussian_wave_at_E(x = np.arange(0,1000), x0 = 200, Ep = 0.1,sigma = 15):
k0 = m/hbar * sqrt(2*(Ep-1/sigma**2/4))
return Gaussian_wave(x = x, x0 = x0, k0 = k0, sigma = sigma)
Let us now check how the two functions above work: let's make a Gaussian_wave and plot it. Note that the Gaussian wave is, in general, a complex number, so we will have to plot two components, real and imaginary.
we will also plot the square of the wavefunction, which tells us the probability of finding a particle at a certain location
x = np.arange(0,1000,1) # make an x coordinate array for the function
wf = Gaussian_wave(x,x0 = 500) # generate in range 0-1000 with step 1, set max position to 500, and use default k0 and sigma
Plot(x,[wf.real,wf.imag,np.abs(wf)**2],size=(10,2),colors = ['blue','red','black'],lw=1)
The kinetic energy matrix is built according to the Finite Difference Method desribed in the Theory section.
$$
\mathbf T = -\frac{\hbar}{2m} \frac {1}{\Delta x^2}\begin{pmatrix}
&-2 &1 &0 &0 &0 &... &(1) \\
&1 &-2 &1 &0 &0 &... &0\\
&0 &1 &-2 &1 &0 &... &0\\
&0 &0 &1 &-2 &1 &... &0\\
&... &... &... &... &... &... &... \\
&(1) &0 &0 &0 &0 &1 &-2 \\
\end{pmatrix}
$$
The (1) in the opposite corner makes T periodic, as described earlier, the x-space will be periodic, when a free particle hits the edge of the $\{x_i\}$ range, it appears on the other side.
Input parameter:
x - the array of x-values corresponding to the range where the wavefunction is calculated. It is only used to find $\Delta x = x_1 - x_0 $ and the number of points (size) of the requested matrix
periodic - if True, the (1) will be added to the matrix for periodicity
Returns: The kinetic energy matrix
def Kinetic_energy(x,periodic = False):
size = len(x) # returns the size (number of points) in x-array
dx = x[1]-x[0] # difference between two nearby points
D2 = (np.diag(-2*np.ones(size)) #make a matrix with its main diagonal having size "size" and fill the diagonal with "-2" values
+ np.diag(np.ones(size-1),1) # make matrix with its diagonal one below the main diagonal, its size will be "size-1", and fill with '1's
+ np.diag(np.ones(size-1),-1)) # diagonal one above the main.
if periodic:
D2[size-1,0] = 1 # these two lines make the T (and Hamiltonian) periodic
D2[0,size-1] = 1 # in temrs of the derivative once it needs a point to the left of 0-point, it assumes it is the last point in the wavefunction
T = (-hbar/m/2) * (1/dx**2) * D2
return T
As described in the Theory section, the potential energy operator is represented by a diagonal matrix :
$$
\mathbf V = \begin{pmatrix}
&V(x_1) &0 &0 &... \\
&0 &V(x_2) &0 &... \\
&0 &0 &V(x_3) &... \\
&... &... &... &... \\
\end{pmatrix}
$$
Input: parameters:
x - the x range to calculate V
position - the position where the barrier starts
height - the energy on the top of the barrier
width - the width of the barrier along x
Returns a list of output objects
V - The potential energy matrix
flatV - the 1D array of $V(x)$, where $V=0$ everywhere except between position and position+width, where its value V =height.
def Potential_barrier(x,position,height,width):
flatV = np.array([height if position< pos < position+width else (0) for pos in x]) #make 1D array
V = np.diag(flatV) #make matrix with given diagonal
return V,flatV
Let's generate a potential barrier and plot it
x = np.arange(0,1000)
V,flatV = Potential_barrier(x,position = 200,height = 0.1,width = 100)
Plot(x,[flatV],size=(10,1))
First, we find the basis functions of the Hamiltonian by solving the time-independent Schrodinger equation in matrix form:
$$
\mathbf{H}\mathbf{\varphi}_i = E_i\mathbf{\varphi}_i
$$
Here, this is achieved by using the numpy function np.linalg.eigh($\mathbf{H}$). This function may take a long time to execute if the size of the input matrix (N, number of points in $\{x_i\}$ array) is too large, i.e., much larger than 1000, as the execution time is proportional to $N^3$. It also has a limit of about N~30,000. One could switch to sparse matrices and respective eigh functions which are designed to work with matrices that have a lot of zero elements, for larger N it may work faster.
The function returns N eigenfunctions $\varphi_i$ and eigenvalues $E_i$.
In the next step, we find the expansion coefficients $c_n$ such that our initial wavepacket is expressed as a linear superposition of basis functions $\varphi_i$: $$ \psi(x,t=0) = \sum_i c_i \varphi_i (x) $$ We will use the following numerical integral to find these coefficients (see the Theory section): $$ c_n = \left<\varphi_n(x) | \psi(x) \right> = \sum_i {\varphi_n^*(x_i) \psi(x_i)}\Delta x = \sum_i {\varphi_{ni}^* \psi_i}\Delta x\\ \mathbf c = \boldsymbol {\varphi^* \psi} \Delta x $$
Input parameters:
psi - the wavefunction for which coefficients need to be found, i.e., an array of $\psi_i = \psi (x_i)$ (complex numbers!)
H - Hamiltonian matrix (T+V), kinetic + potential energy
dx - $\Delta x$
Returns:
cn - The expansion coefficients $c_n$
phi - an array of N eigenvectors (eigenfunctions), or basis states. This is an NxN matrix.
cn_phi - an array of eigenvectors scaled by $c_n$ coefficients, $c_n \varphi_n (x_i)$
E - an array of eigenvalues (the energies of the basis states)
def CalcCn(psi,H,dx):
E, phi = np.linalg.eigh(H)
phi = phi.T # transpose eigenfunctions to get wavefunctions as lines in the matrix, not columns
norm = np.sum((np.abs(phi)**2)*dx,axis = 0) # find sum of squares for all bectors
psi = psi/np.sqrt(norm) # normalize all functions (all you really have to do is divide by sqrt(dx)?)
cn = np.zeros_like(phi[0], dtype=complex)
for j in range(0, len(psi)):
cn[j] = np.sum(np.conj(phi[j]) * psi) *dx # for each eigenvector, compute the inner product
cn_phi = cn * phi.T # multiply basis bunctions with cn coeffs and transpose, this will make the future calculations faster
return cn,phi,cn_phi,E
Below, we will calculate the wavefunction at any time $t$ using the following equation: $$ \psi(x_i,t) = \sum_n c_n \varphi_n (x_i) \exp{\left ( -i E_n t/\hbar \right )} $$
Input:
cn_phi - the set of basis wavefunctions where each is scaled by the expansion coefficients, $c_n \varphi_n (x_i)$
E - the eigenvalues, i.e., the energies of the basis states
t - time at which the wavefunction is calculated
Returns:
The wavefunction at time t
def GetPsi(cn_phi,E,t):
return cn_phi@(np.exp(-1j*E*t)) #note: @ denotes multiplication by matrix rules, not element by element!
# If we would define the arrays as matrices (np.matrix), then * will be by matrix rules
Here, we will generate the initial wavefunction $\psi$, generate the potential and kinetic energy matrices, calculate the wavefunction at 0-time and at time t, and plot the results.
#specify parameters of the problem and plot wavefunction at time = 0 and time=t
t = 800 # time at which to calculate wavefunction in addition to zero time
dx = 1 # step size along x axis
x = np.arange(0,800,dx) # x-array to work in
Ep_set = 0.1 # the target energy of the particle
sigma = 15 # the spread of free particle wavefunction along x
barrier_pos = 400
barrier_height = .2
barrier_width = 3
T = Kinetic_energy(x,periodic=False)
V, flatV = Potential_barrier(x,barrier_pos,barrier_height,barrier_width)
wf = Gaussian_wave_at_E(x,x0 = 300,Ep = Ep_set,sigma=sigma) # initial wavefunction
Ep = (wf.conjugate() @ (T @ wf)).real * dx # check actual expectation value <psi|T|psi> (sligtly differs from Ep_set because of discrete x-axis!)
cn,phi,cn_phi,E = CalcCn(wf,T+V,x[1]-x[0])
print(f'Ep_set = {Ep_set}\nEp calculated = {Ep}\nExpectation energy {np.sum(np.square(np.abs(cn)) * E)}')
maxpop = np.argmax(np.abs(cn)) # state n with maximum population (max cn)
print(f'State with maximum population n={maxpop}')
wft = GetPsi(cn_phi,E,t) # calculate wavefunction at time t
#note: when plotting we will conventionally shift the zero of the wavefunction to the particle's energy level
# we will also plot the |psi|^2, the probability distribution, and the potential
Plot(x,[wf.real+Ep,wf.imag+Ep,np.square(np.abs(wf))*2,flatV],size=(10,2),lw=1) # conventially plot wavefunction with zero at its energy Ep
Plot(x,[wft.real+Ep,wft.imag+Ep,np.square(np.abs(wft))*2,flatV],size=(10,2),lw=1)
Plot(x,[phi[i] for i in range(maxpop-1,maxpop+2)],size=(10,2),lw=.5) # plot some Hamiltonian eigenfunctions
We will calculate and plot the wavefunction at a series of times and assemble these plots into a movie.
Some notes on the program:
Be aware that these calculations may take a long time; the animation will be shown only when all timeframes are calculated.
It is advised to change the time step (last argument in the first function, t = np.arange(0,700,10)) to expedite calculations, in the example here we have to calculate (700-0)/10 time frames.
#specify parameters of the problem
dx = 1
t = np.arange(0,700,10) # time point array for which to calculoate wavefunction
x = np.arange(0,800,dx) # x-array to work in
Ep_set = 0.1 # the energy of the particle
sigma = 15 # bandwidth
x0 = 300 # position
barrier_pos = 400
barrier_height = 0.2
barrier_width = 3
T = Kinetic_energy(x,periodic = True) # make it periodic for smoother spectrum
plt.rcParams['animation.embed_limit'] = 100.0 # set animation limit to 100 MB
V, flatV = Potential_barrier(x,barrier_pos,barrier_height,barrier_width)
wf = Gaussian_wave_at_E(x,x0 = x0,Ep = Ep_set,sigma=sigma)
Ep = (wf.conjugate() @ (T @ wf)).real * dx # make particle energy consistent with expectation value (differs because of discrete x-axis!)
print(f'Ep_set = {Ep_set}, Ep = {Ep}')
cn,phi,cn_phi,E = CalcCn(wf,T+V,x[1]-x[0])
fig = plt.figure(figsize = (10,4) ,dpi=100) # make new plot frame
ax = plt.axes(xlim=(np.min(x),np.max(x)),ylim = (-.4,.3))
ax.set_title(f'Particle E={Ep_set}, Barrier h={barrier_height}, w = {barrier_width}')
ax.plot(x,flatV)
ax.grid()
ax.plot(np.square(np.abs(cn))*2000/dx,E,color = 'black')
maxpop = np.argmax(np.abs(cn)), #find the level that has maximum population
# comment out the following lines if the external LaTex editor is not installed:
if use_latex == True:
plt.rcParams['text.usetex'] = True # this tells print(f'...') to interprete text as LaTeX expression
xxx = np.square(np.abs(cn[maxpop[0]]))*2000/dx
plt.text(xxx,Ep+0.03,'$|c_n(E)|^2$ ' ,horizontalalignment='left', color = 'black',size = 14)
plt.text(xxx,-.4+.05,'$|\psi(x)|^2$ ' ,horizontalalignment='left', color = 'black',size = 14)
plt.text(x[-1],.3-.05,'$\psi_{real}$ .' ,horizontalalignment='right', verticalalignment = 'top',color = 'red',size = 16)
plt.text(x[-1],.3-.1,'$\psi_{imag}$ .' ,horizontalalignment='right', verticalalignment = 'top',color = 'blue',size = 16)
plt.rcParams['text.usetex'] = False
line1, = ax.plot([],[],color = 'red') # reserve lines on plot to show real, imaginary and square of wavefunction
line2, = ax.plot([],[],color = 'blue') # they weill be replotted for each new frame in animation
line3, = ax.plot([],[],color = 'black')
def animate(i): # this function is called by FuncAnimation to make a plot of frame number i
print(f'\rFrame {i+1} of {len(t)}...',end = '')
wft = GetPsi(cn_phi,E,t[i]) # calculate wavefunction for frame i at time t[i]
line1.set_data(x, wft.real + Ep) # change line in current plot
line2.set_data(x, wft.imag + Ep)
line3.set_data(x, np.square(np.abs(wft))*5-.4)
# animation will call 'animate' function 'frame' times and record each 'fig' plot as a new animation frame
ani = FuncAnimation(fig, animate, frames=len(t),interval=50)
plt.close() # prevent from showing plot inline
if save_mp4 == True:
name = f'packet_tunnel_E{Ep_set}_h{barrier_height}_w{barrier_width}.mp4'
# to save plot as mp4 file you need to install ffmpeg codec and att its /bin/ to windows PATH
# if not installed - you will get error, simply comment out the following lines to skip saving to file
print(f'Saving in "{name}"...')
ani.save(name)
print('done')
if show_html == True:
print('Creating HTML animation...')
from IPython.display import HTML
#The following will make animation in shtml language and append HTML code to output below
#if create_HTML:
anim = ani.to_jshtml() #create animation as jshtml string
print('done')
display(HTML(anim)) # for some reason, plain HTML(anim) does not work inside if block... CRAZY!
The approach described above can be used to model wavepacket time evolution in any potential. Here, we set a harmonic potential and animate the wavepacket motion. Note that the k0-vector calculated for the set kinetic energy Ep_set in this case is not the full energy E of the particle, as we ignore the potential energy. In the case of the barrier we created the wavepacket in the area where V=0, now it is not the case.
To begin, it is advised to change the time step (last argument in the first function, t = np.arange(0,1350,10)) to expedite calculations, in the example here we have to calculate (1350-0)/10 timeframes.
#specify parameters of the problem
dx = 1
t = np.arange(0,1350,10) # time point array for which to calculoate wavefunction, use max time 1350 for full period with the parameters below
x = np.arange(0,600,dx) # x-array to work in
Ep_set = 0.1 # the kinetic energy of the particle (does not include the potential energy)
sigma = 15 # wavepacket width
x0 = floor((x[0] + x[-1])/2) # put the wavepaket at the bottom of the potential
width = 300 # this defines the 'width' of the harmonic potential
flatV = (x-x0)**2/width**2 # we create the harmonic potential
V = np.diag(flatV) # and make it a diagonal matrix
T = Kinetic_energy(x)
plt.rcdefaults()
plt.rcParams['animation.embed_limit'] = 100.0 # set animation limit to 100 MB
#plt.rcParams['text.usetex'] = False
wf = Gaussian_wave_at_E(x,x0 = x0,Ep = Ep_set,sigma=sigma)
Ep = (wf.conjugate() @ ((T+V) @ wf)).real * dx # Calculate the total energy of the particle
print(f'Kinetic energy T = {(wf.conjugate() @ (T @ wf)).real * dx}')
print(f'Full energy Ep = T+V = {Ep}')
cn,phi,cn_phi,E = CalcCn(wf,T+V,x[1]-x[0])
maxpop = np.argmax(np.abs(cn)), #find the level that has maximum population:
print(f'Max populated level = {maxpop[0]}, E={E[maxpop[0]]}')
fig = plt.figure(figsize = (10,4) ,dpi=100) # make new plot frame
ax.set_title(f'Wavepacket in harmonic potential')
ax = plt.axes(xlim=(np.min(x),np.max(x)),ylim = (-.3,.4))
ax.plot(x,flatV) # plot potential
xc = width * sqrt(Ep) # this is classical limit for motion, Ep=V
ax.plot([x0+xc,x0+xc],[-.3,.4],color = 'black', lw = 1)
ax.plot([x0-xc,x0-xc],[-.3,.4],color = 'black', lw = 1)
ax.plot([x[0],x[-1]],[0,0],color = 'black', lw = 1.5)
plt.axvspan(x[0], x0-xc, color='gray', alpha=0.3) #plot gray areas (classically forbidden)
plt.axvspan(x0+xc, x[-1], color='gray', alpha=0.3)
ax.plot(np.square(np.abs(cn))*1000/dx,E,color='black') #plot expansion coeffs cn^2 as function of eigen-energies
ax.grid()
plt.text(x[-1],Ep+.02,f'Level n = {maxpop[0]} ',horizontalalignment='right',color = 'blue')
if use_latex == True:
plt.rcParams['text.usetex'] = True
plt.text(x[-1],Ep-.02,'$E = (n+1/2)\hbar \omega$ ',horizontalalignment='right',verticalalignment = 'top', color = 'blue',size=14)
plt.text(np.square(np.abs(cn[maxpop[0]]))*1000/dx,Ep+0.03,'$|c_n (E)|^2$ ' ,horizontalalignment='left', color = 'black',size = 14)
plt.text(np.square(np.abs(cn[maxpop[0]]))*1000/dx,-.3+0.05,'$|\psi (x)|^2$ ' ,horizontalalignment='left', color = 'black',size = 14)
plt.rcParams['text.usetex'] = False
plt.title(f'Particle in harmonic potential: Ep = {Ep:.4f}, V width = {width}')
line1, = ax.plot([],[],color = 'red')
line2, = ax.plot([],[],color = 'blue')
line3, = ax.plot([],[],color = 'black')
def animate(i):
print(f'\rFrame {i+1} of {len(t)} ',end = '')
wft = GetPsi(cn_phi,E,t[i])
line1.set_data(x, wft.real + Ep)
line2.set_data(x, wft.imag + Ep)
line3.set_data(x, np.square(np.abs(wft))*5-.3)
ani = FuncAnimation(fig, animate, frames=len(t),interval=50)
plt.close() # prevent from showing plot inline
if save_mp4 == True:
name = f'packet_harmonic_E{Ep_set}_w{width}.mp4'
print(f'Saving in "{name}"...')
ani.save(name) # need to install ffmpeg codec and add its /bin/ to windows PATH, comment out if not installed
print('done')
if show_html == True:
print('Creating HTML animation...')
from IPython.display import HTML
anim = ani.to_jshtml() #create animation as jshtml string
print('done')
display(HTML(anim)) # for some reason, plain HTML(anim) does not work inside if block... CRAZY!
#print('nMaking HTML video...')
#from IPython.display import HTML
#HTML(ani.to_jshtml())
A wavepacket can also be reflected from an 'antibarrier', i.e., a potential well. You can model that using the potential barrier animation routine by swapping the height to a negative value (barrier_height = - 0.2). Note that the wavefunction reflects from both sides of such a barrier, and, for a certain barrier width, the wave reflected from the front and back sides of the well will interfere destructively, so there is almost no reflected wave! Barrier_height = -0.2, and the wavepacket parameters shown above the parameters:
Ep_set = 0.1 # the energy of the particle
sigma = 15 # bandwidth
x0 = 300 # position
The magic width of the well is 4.5, 9, etc., which is consistent with the phase difference of the two reflected waves being 180 degrees.
The following is the complete Python program that makes the animation of the wavepacket passing the potential barrier. You can copy just that piece and run it.
# By Sergei Savikhin, Purdue University, September 2025
# minimalistic version of the animation of quantum wavepacket crossing the potential barrier
import numpy as np # 'numerical python' library, defines matrices, comlpex numbers etc, we will refer to its functions using prefix 'np'
import matplotlib.pyplot as plt # set of functions for drawing graphs
from matplotlib.animation import FuncAnimation # set of functions to create animations out of plots
from IPython.display import HTML, display # to plot animation to html (inline Jupiter output)
from math import * # input all math functions (sin, cos, exp...), with no need to use any prefix
hbar, m = 1, 1 # set hbar to 1 and mass to 1
def Gaussian_wave_at_E(x = np.arange(0,1000), x0 = 200, Ep = 0.1,sigma = 15):
k0 = m/hbar * sqrt(2*(Ep-1/sigma**2/4))
wf = np.exp( -1/2* (x-x0)**2/sigma**2) *np.exp(1j*k0*x)
return wf/(sqrt(np.sum(np.abs(wf)**2)*(x[1]-x[0]))) # normalize and return
def Kinetic_energy(x,periodic = False):
size = len(x) # returns the size (number of points) in x-array
dx = x[1]-x[0] # difference between two nearby points
D2 = (np.diag(-2*np.ones(size)) #make a matrix with its main diagonal having size "size" and fill the diagonal with "-2" values
+ np.diag(np.ones(size-1),1) # make matrix with its diagonal one below the main diagonal, its size will be "size-1", and fill with '1's
+ np.diag(np.ones(size-1),-1)) # diagonal one above the main.
if periodic:
D2[size-1,0] = 1 # these two lines make the T (and Hamiltonian) periodic
D2[0,size-1] = 1 # in temrs of the derivative once it needs a point to the left of 0-point, it assumes it is the last point in the wavefunction
T = (-hbar/m/2) * (1/dx**2) * D2
return T
def Potential_barrier(x,position,height,width):
flatV = np.array([height if position< pos < position+width else (0) for pos in x]) #make 1D array
V = np.diag(flatV) #make matrix with given diagonal
return V,flatV
def CalcCn(psi,H,dx): # get aigenfunctions of Hamiltonian and express our wavefunction in this basis
E, phi = np.linalg.eigh(H)
phi = phi.T
norm = np.sum((np.abs(phi)**2)*dx,axis = 0) # find sum of squares for all bectors
psi = psi/np.sqrt(norm) # normalize all functions (all you really have to do is divide by sqrt(dx)?)
cn = np.zeros_like(phi[0], dtype=complex)
for j in range(0, len(psi)): cn[j] = np.sum(np.conj(phi[j]) * psi) *dx # for each eigenvector, compute the inner product
cn_phi = cn * phi.T # multiply basis bunctions with cn coeffs and transpose, this will make the future calculations faster
return cn,phi,cn_phi,E
def GetPsi(cn_phi,E,t): # get wavefunction at time t using the eigenfunctions of Hamiltonian
return cn_phi@(np.exp(-1j*E*t)) #note: @ denotes multiplication by matrix rules, not element by element!
#specify parameters of the problem
t = np.arange(0,700,10) # time points array for which to calculoate the wavefunction
x = np.arange(0,800,1) # x-array to work in
Ep, sigma, x0 = 0.1,15,300 # the wavepacket parameters
barrier_pos, barrier_height, barrier_width = 400, 0.2, 3 # barrier parameters
wf = Gaussian_wave_at_E(x,x0 = x0,Ep = Ep,sigma=sigma) # make gaussian wavepacket at time 0
V, flatV = Potential_barrier(x,barrier_pos,barrier_height,barrier_width) # make potential matrix
T = Kinetic_energy(x) # make kinetic energy matrix
cn,phi,cn_phi,E = CalcCn(wf,T+V,x[1]-x[0]) #express the wavefunction in Hamiltonian basis
fig = plt.figure(figsize = (10,4) ,dpi=100) # make new plot frame
ax = plt.axes(xlim=(np.min(x),np.max(x)),ylim = (-.4,.3)) # set axes
ax.plot(x,flatV) # plot potential
line1, = ax.plot([],[],color = 'red') # reserve for psi.real
line2, = ax.plot([],[],color = 'blue') # reserve for psi.imaginary
line3, = ax.plot([],[],color = 'black')# reserve for psi^2
def animate(i):
wft = GetPsi(cn_phi,E,t[i]) # calculate wavefunction at time t[i], i is the animation frame number
line1.set_data(x, wft.real + Ep) # replot lines with new values of psi
line2.set_data(x, wft.imag + Ep)
line3.set_data(x, np.square(np.abs(wft))*5-.4)
ani = FuncAnimation(fig, animate, frames=len(t),interval=50) # call animation recorder. It will call 'animate' with i=frame_number
plt.close() # otherwise the last frame will be displayed as a separate plot
name = f'packet_tunnel_E{Ep}_h{barrier_height}_w{barrier_width}-min.mp4' # set name for mp4 file to record movie to
print(f'Saving in "{name}"...',end = '')
ani.save(name) # need to install ffmpeg codec and att its /bin/ to windows PATH
print('done\nMaking HTML video...')
from IPython.display import HTML
display(HTML(ani.to_jshtml())) # make an HTML movie file and export to browser (in Jupyter Notebook)